import cv2
import os, sys
import glob
import numpy as np
import time
import pdb
import argparse
from pathlib2 import Path

if __name__ == "__main__":

    path_data = Path("/mnt/data/DATA/GRID/processing")

    v_ids = sorted(path_data.glob("*"))
    print(f"Have {len(v_ids)} dir files, \n {v_ids[:10]}")

    # for each speaker
    for i, ipath_id in enumerate(v_ids):
        id_name = ipath_id.stem

        dir_video = path_data.joinpath(id_name)
        v_dir_videos = sorted(dir_video.glob("*/video/video_orignal.mp4"))

        print(f"\nid= {id_name}, Have {len(v_dir_videos)} video files, \n {v_dir_videos[:10]}")

        # for each video
        for ipath_video in v_dir_videos:

            # create data path
            path_save_frames = ipath_video.parent.parent.joinpath("frames")

            # if already exists frames
            if path_save_frames.exists():
                if len(list(path_save_frames.glob("*.jpg"))) > 5:
                    continue

            path_save_frames.mkdir(parents=True, exist_ok=True)
            #print("save frames dir = ", {path_save_frames})

            # STEP1: extract frames
            count = 0
            cap = cv2.VideoCapture(str(ipath_video))
            length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

            while True:

                success, image = cap.read()
                if not success:
                    break

                ipath_frame = path_save_frames.joinpath(f"frame-{count:04d}.jpg")

                if not ipath_frame.exists():
                    cv2.imwrite(str(ipath_frame), image)

                count += 1
            
            #print(f"extract and save the {count}/{length}-th frame --> {ipath_frame}")

            #break
        #break